"""
题目：使用邻接矩阵表示图，实现一个函数计算图中顶点的总数（邻接矩阵的维度即顶点数）。
"""


class GraphMatrix:
    def __init__(self, num_vertices, directed=False):
        self.num_vertices = num_vertices  # 顶点数量（矩阵维度）
        self.directed = directed
        # 初始化邻接矩阵（0表示无边，1表示有边）
        self.matrix = [[0 for _ in range(num_vertices)] for _ in range(num_vertices)]

    def add_edge(self, v1, v2, weight=1):
        if 0 <= v1 < self.num_vertices and 0 <= v2 < self.num_vertices:
            self.matrix[v1][v2] = weight
            if not self.directed:
                self.matrix[v2][v1] = weight
        else:
            raise ValueError("顶点编号超出范围")

    def get_total_vertices(self):
        """返回图中顶点的总数（邻接矩阵的行数/列数）"""
        return self.num_vertices

# 测试
if __name__ == "__main__":
    # 创建3个顶点的无向图
    graph1 = GraphMatrix(num_vertices=3)
    graph1.add_edge(0, 1)
    graph1.add_edge(1, 2)
    print("图1顶点总数:", graph1.get_total_vertices())  # 输出: 3

    # 创建5个顶点的有向图
    graph2 = GraphMatrix(num_vertices=5, directed=True)
    graph2.add_edge(0, 3)
    graph2.add_edge(2, 4)
    print("图2顶点总数:", graph2.get_total_vertices())  # 输出: 5